from .deformable_transformer import Transformer as trans
from torch import nn
from .module.transformerdecoder import TransformerDecoderLayer, TransformerDecoder

def build_transformer(args):
    return Transformer(
        d_model=args.hidden_dim,
        dropout=args.dropout,
        nhead=args.nheads,
        dim_feedforward=args.dim_feedforward,
        num_encoder_layers=args.enc_layers,
        num_dec_layers_hopd=args.dec_layers_hopd,
        num_dec_layers_interaction=args.dec_layers_interaction,
        normalize_before=args.pre_norm,
        return_intermediate_dec=True,
        num_feature_levels=args.num_feature_levels, 
        dec_n_points=args.dec_n_points,  
        enc_n_points=args.enc_n_points,
        two_stage=args.two_stage, 
        two_stage_num_proposals=args.num_queries)

class Transformer(trans):
    def __init__(self, 
                 d_model=512, 
                 nhead=8, 
                 num_encoder_layers=6,
                 num_dec_layers_hopd=3, 
                 num_dec_layers_interaction=3, #new
                 dim_feedforward=2048, 
                 dropout=0.1,
                 activation="relu", 
                 normalize_before=False,
                 return_intermediate_dec=False,
                 num_feature_levels=4, 
                 dec_n_points=4,  
                 enc_n_points=4,
                 two_stage=False, 
                 two_stage_num_proposals=300):
        super(Transformer,self).__init__(d_model,  nhead,  num_encoder_layers, num_dec_layers_hopd, dim_feedforward,  dropout,
                    activation,  normalize_before, return_intermediate_dec, num_feature_levels, dec_n_points, enc_n_points,
                    two_stage, two_stage_num_proposals)
        #d_model, nhead, num_decoder_layers, dim_feedforward, dropout, activation, return_intermediate_dec, num_feature_levels, dec_n_points
        self.human_decoder = self.build_decode(d_model, nhead, num_dec_layers_interaction, dim_feedforward, dropout, activation, return_intermediate_dec, num_feature_levels, dec_n_points)
        self.human_reference_points = nn.Linear(d_model, 2)
        #self.interaction_decoder = self.build_decode_origin(d_model, nhead, num_dec_layers_interaction, dim_feedforward, dropout, activation, normalize_before, return_intermediate_dec)
        self.interaction_decoder = self.build_decode(d_model, nhead, num_dec_layers_interaction, dim_feedforward, dropout, activation, return_intermediate_dec, num_feature_levels, dec_n_points)
        self.interaction_reference_points = nn.Linear(d_model, 2)
        self._reset_parameters()


    def build_decode_origin(self, d_model, nhead, num_decoder_layers, dim_feedforward, dropout, activation, normalize_before, return_intermediate_dec):
        decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout, activation, normalize_before)
        decoder_norm = nn.LayerNorm(d_model)
        decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm, return_intermediate=return_intermediate_dec)
        return decoder

    def decode_forward_origin(self, tgt, memory, mask, query_embed, pos_embed, decoder):
        '''
            query_embed: 100, N, 256  -> tgt: 100, N, 256(一堆0)   
            tgt:    100, N, 256
            memory: 35x35, N, 256 
            mask:   N, 35x35   -> hs:  6, 100, N, 256      
        '''
        hs, self_attn, cross_attn = decoder(tgt, memory, memory_key_padding_mask=mask, query_pos=query_embed, pos=pos_embed)
        if self.return_intermediate_dec:
            hs = hs.transpose(1, 2)
        else:
            hs = hs.unsqueeze(0).transpose(1, 2)
            self_attn = self_attn.unsqueeze(0)
            cross_attn = cross_attn.unsqueeze(0)
        return hs, self_attn, cross_attn
     
        
    def forward(self, srcs, masks, pos_embed, query_embed):
        assert self.two_stage or query_embed is not None
        # prepare input for encoder
        #   reference_points, init_reference_out: N, 100, 2
        #   tgt, query_embed:                     N, 100, 256    
        src, mask, lvl_pos_embed_flatten, spatial_shapes, level_start_index, valid_ratios = \
            self.encode_input_reshape(srcs, masks, pos_embed)
        # encoder  
        # memory : N, WH + WH(35x35), 256
        memory = self.encode_forward(src, mask, spatial_shapes, level_start_index, valid_ratios, lvl_pos_embed_flatten, self.encoder)
        # prepare input for decoder
        #   N,  WH + WH + , 256 
        #   object
        #   reference_points, init_reference_out: N, 100, 2
        #   tgt, query_embed:                     N, 100, 256
        tgt, query_embed, reference_points, init_reference_out, enc_outputs_class, enc_outputs_coord_unact = \
            self.decode_input_reshape(query_embed, memory, mask, spatial_shapes, self.reference_points)
        hs, inter_references = self.decode_forward(tgt, reference_points, memory, mask, query_embed,
                                                  spatial_shapes, level_start_index, valid_ratios, self.decoder)
        #human
        human_query_embed = hs[-1]
        human_tgt, human_query_embed, human_reference_points, human_init_reference_out, enc_outputs_class, enc_outputs_coord_unact = \
            self.decode_input_reshape(human_query_embed, memory, mask, spatial_shapes, self.human_reference_points)
        human_reference_points = human_reference_points + init_reference_out
        human_init_reference_out = human_reference_points
        #   3, N, 100, 2(human_inter_references<1)
        human_decoder_out, human_inter_references = \
            self.decode_forward(human_tgt, human_reference_points, memory, mask, human_query_embed,
                                spatial_shapes, level_start_index, valid_ratios, self.human_decoder)# 6, N, 100, 256  

        #   change
        interaction_query_embed = hs[-1]+ human_decoder_out[-1]
        interaction_tgt, interaction_query_embed, interaction_reference_points, interaction_init_reference_out, enc_outputs_class, enc_outputs_coord_unact = \
            self.decode_input_reshape(interaction_query_embed, memory, mask, spatial_shapes, self.interaction_reference_points)
        
        interaction_reference_points = interaction_reference_points + init_reference_out + human_init_reference_out
        interaction_init_reference_out = interaction_reference_points
        
        interaction_decoder_out, interaction_inter_references = \
            self.decode_forward(interaction_tgt, interaction_reference_points, memory, mask, interaction_query_embed,
                                spatial_shapes, level_start_index, valid_ratios, self.interaction_decoder)# 6, N, 100, 256
        #interaction_decoder_out, interaction_self_attn, interaction_cross_attn = self.decode_forward_origin(interaction_tgt, memory.permute(1,0,2), mask, interaction_query_embed, pos_embed, self.interaction_decoder)

        outputs = {'src':srcs[-1], 'memory':memory, 
                   'out_query':hs, 'init_reference':init_reference_out, 'inter_references':inter_references,
                   'human_decoder_out':human_decoder_out, 'human_init_reference_out':human_init_reference_out, 'human_inter_references':human_inter_references,
                   'interaction_decoder_out':interaction_decoder_out, 'interaction_init_reference':interaction_init_reference_out, 'interaction_inter_references':interaction_inter_references}
        if self.two_stage:
            outputs.append({'enc_outputs_class':enc_outputs_class, 
                            'enc_outputs_coord_unact':enc_outputs_coord_unact})
            return outputs
        return outputs